"""
L-Shaped Domain and Adaptive Refinement

Implements L-shaped domain geometry and adaptive refinement capabilities
for spectral element method simulations. Supports physics-based error
indicators and smooth refinement transitions.
"""

"""
    define_lshape_domain(n_block::Int, dim::Int=2) -> (Array{Int}, Vector{Tuple})

Define L-shaped domain and identify corner points for refinement.
"""
function define_lshape_domain(n_block::Int, dim::Int=2)
    if dim == 2
        return define_lshape_domain_2d(n_block)
    else
        return define_lshape_domain_3d(n_block)
    end
end

"""
    define_lshape_domain_2d(n_block::Int) -> (Array{Int,2}, Vector{Tuple})

Define 2D L-shaped domain.
"""
function define_lshape_domain_2d(n_block::Int)
    # Create L-shaped domain
    matrix_dist = zeros(Int, n_block + 2, n_block + 2)
    
    # Mark the L-shape with 1s (using the pattern from Python)
    matrix_dist[2:3, 2:3] .= 1
    matrix_dist[4, 3] = 1
    
    # Identify corner point where refinement is needed (at coordinates 2,2)
    corner_points = [(2, 2)]
    
    return matrix_dist, corner_points
end

"""
    define_lshape_domain_3d(n_block::Int) -> (Array{Int,3}, Vector{Tuple})

Define 3D L-shaped domain.
"""
function define_lshape_domain_3d(n_block::Int)
    # Create empty 3D domain
    matrix_dist = zeros(Int, n_block + 2, n_block + 2, n_block + 2)
    
    # Mark the L-shape with 1s - extending the 2D L-shape into 3D
    matrix_dist[2:3, 2:3, 2:3] .= 1
    matrix_dist[4, 3, 2:3] .= 1
    matrix_dist[2:3, 4, 3] .= 1
    
    # Define corner point (interior corner of the L)
    corner_points = [(2, 2, 2)]
    
    return matrix_dist, corner_points
end

"""
    create_refinement_map(matrix_dist::Array{Int}, corner_points::Vector{Tuple}, 
                         levels::Int, dim::Int) -> Array{Int}

Create a multi-level refinement map centered at corner points.
"""
function create_refinement_map(matrix_dist::AbstractArray{Int}, corner_points::AbstractVector{<:Tuple}, 
                              levels::Int, dim::Int)
    refinement_map = zeros(Int, size(matrix_dist))
    
    # For each corner point, apply multi-level refinement
    for corner in corner_points
        if dim == 2
            i, j = corner
            apply_2d_refinement!(refinement_map, matrix_dist, i, j, levels)
        else
            i, j, k = corner
            apply_3d_refinement!(refinement_map, matrix_dist, i, j, k, levels)
        end
    end
    
    return refinement_map
end

"""
    apply_2d_refinement!(refinement_map::Array{Int,2}, matrix_dist::Array{Int,2}, 
                        i::Int, j::Int, levels::Int)

Apply 2D refinement pattern around a corner point.
"""
function apply_2d_refinement!(refinement_map::Array{Int,2}, matrix_dist::Array{Int,2}, 
                             i::Int, j::Int, levels::Int)
    # Highest refinement at corner
    refinement_map[i, j] = levels
    
    # Gradually decrease refinement level in surrounding cells
    for level in 1:(levels-1)
        radius = level
        for di in -radius:radius
            for dj in -radius:radius
                # Only modify points that are in the domain (value 1 in matrix_dist)
                ni, nj = i + di, j + dj
                if (1 <= ni <= size(refinement_map, 1) && 
                    1 <= nj <= size(refinement_map, 2) &&
                    matrix_dist[ni, nj] == 1)
                    # Only set if current value is lower (to maintain highest near corner)
                    current_level = levels - level
                    if refinement_map[ni, nj] < current_level
                        refinement_map[ni, nj] = current_level
                    end
                end
            end
        end
    end
end

"""
    apply_3d_refinement!(refinement_map::Array{Int,3}, matrix_dist::Array{Int,3}, 
                        i::Int, j::Int, k::Int, levels::Int)

Apply 3D refinement pattern around a corner point.
"""
function apply_3d_refinement!(refinement_map::Array{Int,3}, matrix_dist::Array{Int,3}, 
                             i::Int, j::Int, k::Int, levels::Int)
    # Highest refinement at corner
    refinement_map[i, j, k] = levels
    
    # Gradually decrease refinement level in surrounding cells
    for level in 1:(levels-1)
        radius = level
        for di in -radius:radius
            for dj in -radius:radius
                for dk in -radius:radius
                    # Only modify points that are in the domain (value 1 in matrix_dist)
                    ni, nj, nk = i + di, j + dj, k + dk
                    if (1 <= ni <= size(refinement_map, 1) && 
                        1 <= nj <= size(refinement_map, 2) &&
                        1 <= nk <= size(refinement_map, 3) &&
                        matrix_dist[ni, nj, nk] == 1)
                        # Only set if current value is lower (to maintain highest near corner)
                        current_level = levels - level
                        if refinement_map[ni, nj, nk] < current_level
                            refinement_map[ni, nj, nk] = current_level
                        end
                    end
                end
            end
        end
    end
end

"""
    calculate_physics_error_indicators(u::Array, v::Array, w::Union{Array,Nothing}, 
                                      D1::Array, D2::Array, nu::Float64, 
                                      dim::Int) -> Array

Calculate multi-criteria physics-based error indicators for adaptive refinement.
"""
function calculate_physics_error_indicators(u::Array, v::Array, w::Union{Array,Nothing}, 
                                           D1::Array, D2::Array, nu::Float64, 
                                           dim::Int)
    # Initialize unified error indicator
    err_indicator = zeros(eltype(u), size(u))
    
    if dim == 2
        return calculate_2d_error_indicators(u, v, D1, D2, nu)
    else
        return calculate_3d_error_indicators(u, v, w, D1, D2, nu)
    end
end

"""
    calculate_2d_error_indicators(u::Array, v::Array, D1::Array, D2::Array, nu::Float64) -> Array

Calculate 2D physics-based error indicators.
"""
function calculate_2d_error_indicators(u::Array, v::Array, D1::Array, D2::Array, nu::Float64)
    # 1. Vorticity-based indicator (curl of velocity)
    dudy = compute_tensor_derivative(u, D1, 1)  # ∂u/∂y
    dvdx = compute_tensor_derivative(v, D1, 0)  # ∂v/∂x
    
    # 2D vorticity (scalar)
    vorticity = dvdx - dudy
    vorticity_mag = abs.(vorticity)
    
    # Vorticity gradient magnitude
    dvorticity_dx = compute_tensor_derivative(vorticity, D1, 0)
    dvorticity_dy = compute_tensor_derivative(vorticity, D1, 1)
    vort_grad_mag = sqrt.(dvorticity_dx.^2 + dvorticity_dy.^2)
    
    # 2. Velocity gradient indicator (captures shear layers)
    dudx = compute_tensor_derivative(u, D1, 0)  # ∂u/∂x
    dvdy = compute_tensor_derivative(v, D1, 1)  # ∂v/∂y
    
    # Velocity gradient magnitude (Frobenius norm)
    vel_grad_mag = sqrt.(dudx.^2 + dudy.^2 + dvdx.^2 + dvdy.^2)
    
    # 3. Q-criterion (second invariant of velocity gradient tensor)
    Q = 0.5 * ((dudx.^2 + dvdy.^2) - (dudy .* dvdx))
    
    # 4. Combine indicators with weighting
    vort_norm = normalize_indicator(vorticity_mag)
    vort_grad_norm = normalize_indicator(vort_grad_mag)
    vel_grad_norm = normalize_indicator(vel_grad_mag)
    
    # Weight more heavily toward vorticity gradient for corner singularities
    combined_indicator = 0.2 * vort_norm + 0.5 * vort_grad_norm + 0.3 * vel_grad_norm
    
    # 5. Apply Reynolds number scaling
    Re = 1.0 / nu
    reynolds_factor = log10(max(Re, 10)) / 3.0
    
    # Scale indicator by Reynolds number factor
    err_indicator = combined_indicator * reynolds_factor
    
    return err_indicator
end

"""
    calculate_3d_error_indicators(u::Array, v::Array, w::Array, D1::Array, D2::Array, nu::Float64) -> Array

Calculate 3D physics-based error indicators.
"""
function calculate_3d_error_indicators(u::Array, v::Array, w::Array, D1::Array, D2::Array, nu::Float64)
    # 1. 3D Vorticity components
    dudz = compute_tensor_derivative(u, D1, 2)  # ∂u/∂z
    dvdz = compute_tensor_derivative(v, D1, 2)  # ∂v/∂z
    dwdx = compute_tensor_derivative(w, D1, 0)  # ∂w/∂x
    dwdy = compute_tensor_derivative(w, D1, 1)  # ∂w/∂y
    dudy = compute_tensor_derivative(u, D1, 1)  # ∂u/∂y
    dvdx = compute_tensor_derivative(v, D1, 0)  # ∂v/∂x
    
    # 3D vorticity vector (ωx, ωy, ωz)
    vort_x = dwdy - dvdz
    vort_y = dudz - dwdx
    vort_z = dvdx - dudy
    
    # Vorticity magnitude
    vorticity_mag = sqrt.(vort_x.^2 + vort_y.^2 + vort_z.^2)
    
    # 2. Velocity gradient magnitude
    dudx = compute_tensor_derivative(u, D1, 0)
    dvdy = compute_tensor_derivative(v, D1, 1)
    dwdz = compute_tensor_derivative(w, D1, 2)
    
    vel_grad_mag = sqrt.(dudx.^2 + dudy.^2 + dudz.^2 +
                        dvdx.^2 + dvdy.^2 + dvdz.^2 +
                        dwdx.^2 + dwdy.^2 + dwdz.^2)
    
    # 3. 3D Q-criterion
    Q = 0.5 * ((dudx.^2 + dvdy.^2 + dwdz.^2) - 
               (dudx .* dvdy + dudx .* dwdz + dvdy .* dwdz) -
               (dudy .* dvdx + dudz .* dwdx + dvdz .* dwdy))
    
    # 4. Combine indicators
    vort_norm = normalize_indicator(vorticity_mag)
    vel_grad_norm = normalize_indicator(vel_grad_mag)
    
    combined_indicator = 0.4 * vort_norm + 0.6 * vel_grad_norm
    
    # 5. Reynolds number scaling
    Re = 1.0 / nu
    reynolds_factor = log10(max(Re, 10)) / 3.0
    
    err_indicator = combined_indicator * reynolds_factor
    
    return err_indicator
end

"""
    compute_tensor_derivative(field::Array, D1::Array, direction::Int) -> Array

Compute tensor derivative along specified direction using tensordot equivalent.
"""
function compute_tensor_derivative(field::Array, D1::Array, direction::Int)
    # Julia equivalent of np.tensordot operations from Python
    ndims_field = ndims(field)
    
    if ndims_field == 2
        if direction == 0  # x-direction
            return transpose(D1 * transpose(field))
        else  # y-direction (direction == 1)
            return D1 * field
        end
    elseif ndims_field == 3
        if direction == 0  # x-direction
            result = zeros(eltype(field), size(field))
            for k in 1:size(field, 3)
                for j in 1:size(field, 2)
                    result[:, j, k] = D1 * field[:, j, k]
                end
            end
            return result
        elseif direction == 1  # y-direction
            result = zeros(eltype(field), size(field))
            for k in 1:size(field, 3)
                for i in 1:size(field, 1)
                    result[i, :, k] = D1 * field[i, :, k]
                end
            end
            return result
        else  # z-direction (direction == 2)
            result = zeros(eltype(field), size(field))
            for j in 1:size(field, 2)
                for i in 1:size(field, 1)
                    result[i, j, :] = D1 * field[i, j, :]
                end
            end
            return result
        end
    else
        error("Unsupported field dimensionality: $ndims_field")
    end
end

"""
    normalize_indicator(indicator::Array) -> Array

Normalize error indicator to [0,1] range.
"""
function normalize_indicator(indicator::Array)
    max_val = maximum(indicator)
    return max_val > 0 ? indicator / max_val : indicator
end

"""
    smooth_refinement_transitions!(refinement_map::Array{Int}) -> Array{Int}

Ensure smooth transitions between refinement levels (no jumps > 1).
"""
function smooth_refinement_transitions!(refinement_map::Array{Int})
    smoothed_map = copy(refinement_map)
    dims = size(refinement_map)
    
    # Iterate until no more changes are needed
    while true
        changes_made = false
        
        if length(dims) == 2
            changes_made = smooth_2d_transitions!(smoothed_map, refinement_map)
        else
            changes_made = smooth_3d_transitions!(smoothed_map, refinement_map)
        end
        
        # If no changes were made, we're done
        if !changes_made
            break
        end
        
        # Update refinement_map for next iteration
        refinement_map .= smoothed_map
    end
    
    return smoothed_map
end

"""
    smooth_2d_transitions!(smoothed_map::Array{Int,2}, refinement_map::Array{Int,2}) -> Bool

Smooth transitions in 2D refinement map.
"""
function smooth_2d_transitions!(smoothed_map::Array{Int,2}, refinement_map::Array{Int,2})
    n_rows, n_cols = size(refinement_map)
    changes_made = false
    
    for i in 1:n_rows
        for j in 1:n_cols
            if refinement_map[i, j] > 0  # Only check cells in the domain
                # Check all neighbors
                for (di, dj) in [(-1,0), (1,0), (0,-1), (0,1)]
                    ni, nj = i + di, j + dj
                    if (1 <= ni <= n_rows && 1 <= nj <= n_cols && 
                        refinement_map[ni, nj] > 0)  # Only check domain cells
                        # If difference is more than 1, adjust
                        if refinement_map[i, j] > refinement_map[ni, nj] + 1
                            smoothed_map[ni, nj] = refinement_map[i, j] - 1
                            changes_made = true
                        end
                    end
                end
            end
        end
    end
    
    return changes_made
end

"""
    smooth_3d_transitions!(smoothed_map::Array{Int,3}, refinement_map::Array{Int,3}) -> Bool

Smooth transitions in 3D refinement map.
"""
function smooth_3d_transitions!(smoothed_map::Array{Int,3}, refinement_map::Array{Int,3})
    n_rows, n_cols, n_depth = size(refinement_map)
    changes_made = false
    
    for i in 1:n_rows
        for j in 1:n_cols
            for k in 1:n_depth
                if refinement_map[i, j, k] > 0  # Only check cells in the domain
                    # Check all neighbors
                    for (di, dj, dk) in [(0,0,1), (0,0,-1), (0,1,0), (0,-1,0), (1,0,0), (-1,0,0)]
                        ni, nj, nk = i + di, j + dj, k + dk
                        if (1 <= ni <= n_rows && 1 <= nj <= n_cols && 
                            1 <= nk <= n_depth && refinement_map[ni, nj, nk] > 0)
                            # If difference is more than 1, adjust
                            if refinement_map[i, j, k] > refinement_map[ni, nj, nk] + 1
                                smoothed_map[ni, nj, nk] = refinement_map[i, j, k] - 1
                                changes_made = true
                            end
                        end
                    end
                end
            end
        end
    end
    
    return changes_made
end

"""
    adaptive_refinement_iteration!(multidomain::MultiDomain{D}, u::Array, v::Array, 
                                  w::Union{Array,Nothing}, D1::Array, D2::Array, 
                                  nu::Float64, tolerance::Float64=0.014) where D

Run one iteration of adaptive refinement with physics-based error indicators.
"""
function adaptive_refinement_iteration!(multidomain::MultiDomain{D}, u::Array, v::Array, 
                                       w::Union{Array,Nothing}, D1::Array, D2::Array, 
                                       nu::Float64, tolerance::Float64=0.014) where D
    # Calculate advanced physics-based error indicators
    error = calculate_physics_error_indicators(u, v, w, D1, D2, nu, D)
    
    # Update refinement map based on error indicators
    if multidomain.refinement_map !== nothing
        update_refinement_from_error!(multidomain.refinement_map, multidomain.matrix_dist, 
                                    error, tolerance, multidomain.corner_points)
        
        # Ensure smooth transitions
        smooth_refinement_transitions!(multidomain.refinement_map)
    end
    
    return error
end

"""
    update_refinement_from_error!(refinement_map::Array{Int}, matrix_dist::Array{Int},
                                 error::Array, tolerance::Float64, corner_points::Vector{Tuple})

Update refinement map based on error indicators.
"""
function update_refinement_from_error!(refinement_map::Array{Int}, matrix_dist::Array{Int},
                                      error::Array, tolerance::Float64, corner_points::Vector{Tuple})
    # Normalize error indicators
    error_max = maximum(error)
    error_normalized = error_max > 0 ? error / error_max : error
    
    # Define error thresholds for refinement levels
    thresh_high = tolerance
    thresh_medium = tolerance * 0.5
    thresh_low = tolerance * 0.25
    max_level = maximum(refinement_map[refinement_map .> 0])
    
    # Update refinement based on normalized error indicators
    for idx in CartesianIndices(matrix_dist)
        if matrix_dist[idx] == 1  # Only consider domain cells
            if Tuple(idx) ∉ corner_points  # Skip corner points (handled separately)
                if error_normalized[idx] > thresh_high
                    # High error: maximum refinement
                    refinement_map[idx] = max_level
                elseif error_normalized[idx] > thresh_medium
                    # Medium error: medium refinement
                    refinement_map[idx] = max(refinement_map[idx], max_level - 1)
                elseif error_normalized[idx] > thresh_low
                    # Low error: slight refinement
                    refinement_map[idx] = max(refinement_map[idx], max_level - 2)
                end
                # Otherwise, keep current refinement level
            end
        end
    end
end

"""
    get_polynomial_order_from_refinement(refinement_level::Int, base_p::Int=2, max_p::Int=5) -> Int

Map refinement levels to polynomial orders.
"""
function get_polynomial_order_from_refinement(refinement_level::Int, base_p::Int=2, max_p::Int=5)
    if refinement_level <= 0
        return base_p
    end
    
    # Linear mapping from refinement level to polynomial order
    p_order = base_p + refinement_level * (max_p - base_p) ÷ 3
    return min(p_order, max_p)
end

"""
    create_polynomial_order_map(refinement_map::Array{Int}, matrix_dist::Array{Int}) -> Array{Int}

Generate a map of polynomial orders from refinement levels.
"""
function create_polynomial_order_map(refinement_map::Array{Int}, matrix_dist::Array{Int})
    p_order_map = zeros(Int, size(refinement_map))
    
    for idx in CartesianIndices(refinement_map)
        if matrix_dist[idx] == 1  # Only process domain cells
            p_order_map[idx] = get_polynomial_order_from_refinement(refinement_map[idx])
        end
    end
    
    return p_order_map
end